monai/networks/layers/filtering.py (1)

406-430: ⚠️ Potential issue | 🟠 Major

TrainableJointBilateralFilter.forward() not updated with the same fix.

This method still uses len_input directly instead of computing spatial_dims = len_input - 2. It will reject valid 3D inputs just like the original bug in TrainableBilateralFilter. Also missing the minimum dimension validation added to the other class.

Proposed fix

     def forward(self, input_tensor, guidance_tensor):
+        if len(input_tensor.shape) < 3:
+            raise ValueError(
+                f"Input must have at least 3 dimensions (batch, channel, *spatial_dims), got {len(input_tensor.shape)}"
+            )
         if input_tensor.shape[1] != 1:
             raise ValueError(
                 f"Currently channel dimensions >1 ({input_tensor.shape[1]}) are not supported. "
                 "Please use multiple parallel filter layers if you want "
                 "to filter multiple channels."
             )
         if input_tensor.shape != guidance_tensor.shape:
             raise ValueError(
                 "Shape of input image must equal shape of guidance image."
                 f"Got {input_tensor.shape} and {guidance_tensor.shape}."
             )

         len_input = len(input_tensor.shape)
+        spatial_dims = len_input - 2

         # C++ extension so far only supports 5-dim inputs.
-        if len_input == 3:
+        if spatial_dims == 1:
             input_tensor = input_tensor.unsqueeze(3).unsqueeze(4)
             guidance_tensor = guidance_tensor.unsqueeze(3).unsqueeze(4)
-        elif len_input == 4:
+        elif spatial_dims == 2:
             input_tensor = input_tensor.unsqueeze(4)
             guidance_tensor = guidance_tensor.unsqueeze(4)

-        if self.len_spatial_sigma != len_input:
-            raise ValueError(f"Spatial dimension ({len_input}) must match initialized len(spatial_sigma).")
+        if self.len_spatial_sigma != spatial_dims:
+            raise ValueError(f"Spatial dimension ({spatial_dims}) must match initialized len(spatial_sigma).")

         prediction = TrainableJointBilateralFilterFunction.apply(
             input_tensor, guidance_tensor, self.sigma_x, self.sigma_y, self.sigma_z, self.sigma_color
         )

         # Make sure to return tensor of the same shape as the input.
-        if len_input == 3:
+        if spatial_dims == 1:
             prediction = prediction.squeeze(4).squeeze(3)
-        elif len_input == 4:
+        elif spatial_dims == 2:
             prediction = prediction.squeeze(4)

         return prediction

monai/networks/layers/filtering.py (1)

233-267: ⚠️ Potential issue | 🟠 Major

Add module-wrapper regression tests for the new validation logic.

Current referenced tests exercise TrainableBilateralFilterFunction.apply() directly, while the changed logic is in TrainableBilateralFilter.forward() and TrainableJointBilateralFilter.forward(). Please add tests for: valid 1D/2D/3D wrapper inputs, ndim < 3, len(spatial_sigma) mismatch, and joint input/guidance shape mismatch.

Suggested test additions

+ # tests/networks/layers/filtering/test_trainable_bilateral.py
+ def test_trainable_bilateral_wrapper_accepts_3d_shape():
+     layer = TrainableBilateralFilter(spatial_sigma=(1.0, 1.0, 1.0), color_sigma=0.2)
+     x = torch.randn(1, 1, 10, 10, 10, dtype=torch.double)
+     y = layer(x)
+     assert y.shape == x.shape
+
+ def test_trainable_bilateral_wrapper_rejects_rank_lt_3():
+     layer = TrainableBilateralFilter(spatial_sigma=(1.0,), color_sigma=0.2)
+     with pytest.raises(ValueError):
+         layer(torch.randn(1, 1))
+
+ # tests/networks/layers/filtering/test_trainable_joint_bilateral.py
+ def test_trainable_joint_wrapper_rejects_shape_mismatch():
+     layer = TrainableJointBilateralFilter(spatial_sigma=(1.0, 1.0, 1.0), color_sigma=0.2)
+     x = torch.randn(1, 1, 10, 10, 10, dtype=torch.double)
+     g = torch.randn(1, 1, 10, 10, 9, dtype=torch.double)
+     with pytest.raises(ValueError):
+         layer(x, g)

As per coding guidelines, "Ensure new or modified definitions will be covered by existing or new unit tests."

Also applies to: 406-447

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@monai/networks/layers/filtering.py` around lines 233 - 267, Add module-level
unit tests that call the wrapper classes (use module(input) or
TrainableBilateralFilter.forward/TrainableJointBilateralFilter.forward) instead
of invoking TrainableBilateralFilterFunction.apply() directly: (1) validate
successful forward passes for valid 1D, 2D, and 3D inputs (batch, channel=1,
spatial dims) matching len(spatial_sigma); (2) assert a ValueError is raised
when ndim < 3; (3) assert a ValueError is raised when len(spatial_sigma) does
not match spatial dims (triggering the check in
TrainableBilateralFilter.forward); and (4) for
TrainableJointBilateralFilter.forward add tests that assert a shape-mismatch
between input and guidance tensors raises the expected error. Ensure tests
construct modules with differing len_spatial_sigma and call the module (not the
C++ function) so the new validation logic in TrainableBilateralFilter.forward
and TrainableJointBilateralFilter.forward is covered.

monai/networks/layers/filtering.py (1)

233-267: ⚠️ Potential issue | 🟠 Major

Add regression tests for the new forward validation branches.

This PR changes validation/control flow in both forward methods (Line 234 onward and Line 407 onward), but no test updates are shown here. Please add or point to tests covering: valid 1D/2D/3D inputs, <3 rank rejection, and len_spatial_sigma mismatch rejection in both classes.

As per coding guidelines, "Ensure new or modified definitions will be covered by existing or new unit tests."

Also applies to: 406-447

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@monai/networks/layers/filtering.py` around lines 233 - 267, The forward()
validation branches now added need unit tests: add tests exercising
TrainableBilateralFilter.forward and the other class's forward (the second
forward around lines 406-447) for valid 1D/2D/3D inputs returning same shaped
tensors, for inputs with rank < 3 raising ValueError, for inputs with channel
dimension != 1 raising ValueError, and for cases where spatial_dims !=
self.len_spatial_sigma raising ValueError; implement tests by constructing small
tensors of appropriate shapes, calling the respective forward methods (or the
module forward via model(input_tensor)), and asserting output shapes or that
ValueError is raised, referencing the methods forward,
TrainableBilateralFilterFunction.apply, and the attribute len_spatial_sigma to
locate code under test.